package simpletrace

import (
	"context"
	"fmt"
)

type RedisClient interface {
	MSetExWithCtx(ctx context.Context, keyValues map[string]any, expire int64) error
	MGetWithCtx(ctx context.Context, keys ...string) ([]any, error)
	SAddWithCtx(context.Context, string, ...any) (int64, error)
	ExpireWithCtx(context.Context, string, int64) (int64, error)
	GetWithCtx(context.Context, string) (any, error)
	SetWithCtx(context.Context, string, any) error
	SetExWithCtx(context.Context, string, any, int64) error
	String(any, error) (string, error)
	Bytes(any, error) ([]byte, error)
	DelWithCtx(context.Context, ...string) (int64, error)
	SMembersWithCtx(context.Context, string) ([]string, error)
	SCardWithCtx(context.Context, string) (int64, error)
	ExistsWithCtx(context.Context, string) (bool, error)
	IncrByWithCtx(context.Context, string, int64) (int64, error)
}

type RedisExporterOption struct {
	TaskSize        int
	MaxTaskSpanSize int
	SetTraceIDFn    SetTraceIDFn
	OnTraceCreated  OnTraceExportedFn
	SetTraceSpansFn SetTraceSpansFn
	KeyPrefix       string
	ExpireTime      int64
}

type RedisExporterOptionFn func(*RedisExporterOption)

func WithRETaskSize(size int) RedisExporterOptionFn {
	return func(o *RedisExporterOption) {
		o.TaskSize = size
	}
}

func WithREMaxTaskSpanSize(size int) RedisExporterOptionFn {
	return func(o *RedisExporterOption) {
		o.MaxTaskSpanSize = size
	}
}

func WithRESetTraceIDFn(redisCli RedisClient) RedisExporterOptionFn {
	return func(o *RedisExporterOption) {
		o.SetTraceIDFn = setTraceIDByRedis(redisCli)
	}
}

func WithRESetTraceSpansFn(redisCli RedisClient) RedisExporterOptionFn {
	return func(o *RedisExporterOption) {
		o.SetTraceSpansFn = setSpansByRedis(redisCli)
	}
}

func WithREOnTraceExportedFn(fn OnTraceExportedFn) RedisExporterOptionFn {
	return func(o *RedisExporterOption) {
		o.OnTraceCreated = fn
	}
}

func WithREKeyPrefix(prefix string) RedisExporterOptionFn {
	return func(o *RedisExporterOption) {
		o.KeyPrefix = prefix
	}
}

func WithREExpireTime(expire int64) RedisExporterOptionFn {
	return func(o *RedisExporterOption) {
		o.ExpireTime = expire
	}
}

func NewRedisExporter(opts ...RedisExporterOptionFn) Exporter {
	param := &RedisExporterOption{}
	for _, opt := range opts {
		opt(param)
	}
	if param.KeyPrefix != "" {
		withRedisKeyPrefx(param.KeyPrefix)
	}
	if param.ExpireTime > 0 {
		withRedisExpireTime(param.ExpireTime)
	}
	return NewDefaultExporter(
		WithDefTaskSize(param.TaskSize),
		WithDefMaxTaskSpanSize(param.MaxTaskSpanSize),
		WithDefSetTraceIDFn(param.SetTraceIDFn),
		WithDefSetTraceSpans(param.SetTraceSpansFn),
		WithDefOnTraceExportedFn(param.OnTraceCreated),
	)
}

var (
	RedisKeyTraceSpanInfo   = "simpletrace:span_info:%s:%s"           // %s for traceID and spanID
	RedisKeyTraceSpanIDSet  = "simpletrace:span_id_set:%s"            // %s for traceID
	RedisKeyTraceIDSet      = "simpletrace:trace_id_set"              // Set of all trace IDs
	RedisKeyTraceRootSpanID = "simpletrace:root_span_id:%s"           // %s for traceID
	RedusKeyTraceMemUsage   = "simpletrace:stat_data:memory_usage:%s" // %s for traceID

	DefaultExpireTime = int64(900) // 10 min
)

func withRedisKeyPrefx(keyPrefix string) {
	RedisKeyTraceSpanInfo = keyPrefix + RedisKeyTraceSpanInfo
	RedisKeyTraceSpanIDSet = keyPrefix + RedisKeyTraceSpanIDSet
	RedisKeyTraceIDSet = keyPrefix + RedisKeyTraceIDSet
	RedisKeyTraceRootSpanID = keyPrefix + RedisKeyTraceRootSpanID
	RedusKeyTraceMemUsage = keyPrefix + RedusKeyTraceMemUsage
}

func withRedisExpireTime(expire int64) {
	DefaultExpireTime = expire
}

func setTraceIDByRedis(redisClient RedisClient) func(ctx context.Context, traceID, spanID string) error {
	return func(ctx context.Context, traceID, spanID string) error {
		isExist := false
		if exists, err := redisClient.ExistsWithCtx(ctx, RedisKeyTraceIDSet); err != nil {
			logger.ErrorCtx(ctx, "Failed to check if trace ID set exists: %v", err)
		} else {
			isExist = exists
		}
		if _, err := redisClient.SAddWithCtx(ctx, RedisKeyTraceIDSet, traceID); err != nil {
			logger.ErrorCtx(ctx, "Failed to add trace ID to set: %v", err)
			return err
		}
		if !isExist {
			if _, err := redisClient.ExpireWithCtx(ctx, RedisKeyTraceIDSet, DefaultExpireTime); err != nil {
				logger.ErrorCtx(ctx, "Failed to set expiration for trace ID set: %v", err)
				return err
			}
		}
		setKey := fmt.Sprintf(RedisKeyTraceSpanIDSet, traceID)
		if _, err := redisClient.SAddWithCtx(ctx, setKey, spanID); err != nil {
			logger.ErrorCtx(ctx, "Failed to add span ID to set: %v", err)
			return err
		}
		if _, err := redisClient.ExpireWithCtx(ctx, setKey, DefaultExpireTime); err != nil {
			logger.ErrorCtx(ctx, "Failed to set expiration for span ID set: %v", err)
			return err
		}
		return nil
	}
}

func setSpansByRedis(redisClient RedisClient) func(ctx context.Context, spans []*Span) error {
	return func(ctx context.Context, spans []*Span) error {
		// Prepare data for Redis
		keyValues := make(map[string]any)
		traceMems := make(map[string]int64)
		for _, span := range spans {
			spanStr, err := span.String()
			if err != nil {
				logger.ErrorCtx(ctx, "Failed to convert span to string: %v", err)
				continue
			}
			key := fmt.Sprintf(RedisKeyTraceSpanInfo, span.TraceID, span.SpanID)
			keyValues[key] = spanStr // Assuming span can be serialized to a string or JSON

			// Add span ID to the set for the trace
			setKey := fmt.Sprintf(RedisKeyTraceSpanIDSet, span.TraceID)
			if _, err := redisClient.SAddWithCtx(ctx, setKey, span.SpanID); err != nil {
				logger.ErrorCtx(ctx, "Failed to add span ID to set: %v", err)
				continue
			}

			if span.SpanID == "1" && span.TraceID != "" {
				logger.WarnCtx(ctx, "Set root span ID for trace %s: %s", span.TraceID, span.SpanID)
				// If it's the root span, set the root span ID
				rootSpanKey := fmt.Sprintf(RedisKeyTraceRootSpanID, span.TraceID)
				if err := redisClient.SetWithCtx(ctx, rootSpanKey, spanStr); err != nil {
					logger.ErrorCtx(ctx, "Failed to set root span ID: %v", err)
					continue
				}
				if _, err := redisClient.ExpireWithCtx(ctx, rootSpanKey, DefaultExpireTime); err != nil {
					logger.ErrorCtx(ctx, "Failed to set expiration for root span ID: %v", err)
					continue
				}
				logger.WarnCtx(ctx, "Set root span ID for trace success %s: %s", span.TraceID, span.SpanID)
			}

			traceMems[span.TraceID] += int64(len(spanStr))
		}

		// Store spans in Redis
		err := redisClient.MSetExWithCtx(ctx, keyValues, DefaultExpireTime)
		if err != nil {
			logger.ErrorCtx(ctx, "Failed to store spans in Redis: %v", err)
			return err
		}

		// Store memory usage for each trace
		for traceID, memUsage := range traceMems {
			memKey := fmt.Sprintf(RedusKeyTraceMemUsage, traceID)
			if _, err := redisClient.IncrByWithCtx(ctx, memKey, memUsage); err != nil {
				logger.ErrorCtx(ctx, "Failed to set memory usage for trace %s: %v", traceID, err)
				continue
			}
			if _, err := redisClient.ExpireWithCtx(ctx, memKey, DefaultExpireTime); err != nil {
				logger.ErrorCtx(ctx, "Failed to set expiration for memory usage key %s: %v", memKey, err)
				continue
			}
		}

		return nil
	}
}
